{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "满足以下要求：\n",
    "\n",
    "    同时读取可见光（visible）和红外（infrared）两张图像，通道级联后作为模型输入（4 通道 RGB+IR）。\n",
    "    读取 YOLO 格式的标签（class cx cy w h，归一化）。\n",
    "    支持训练 / 验证（test）划分。\n",
    "    支持在线数据增强（Albumentations）。\n",
    "    返回 dict，方便后续模型直接 inputs = batch['image']，targets = batch['label']。\n",
    "\n",
    "补充说明\n",
    "\n",
    "    输入通道\n",
    "    如果只想用红外或可见光，把 torch.cat([vis, ir], dim=0) 改成单一通道即可。\n",
    "    标签格式\n",
    "    输出的 label 已转换为 [cls, cx, cy, w, h]；若后续模型需 [cx, cy, w, h, cls] 可自行调整顺序。\n",
    "    空标签\n",
    "    无目标图片返回 targets.shape=(0,5)，训练时需做过滤或特殊处理。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "' \\nimport albumentations as A\\n# pip install albumentations\\nfrom albumentations.pytorch import ToTensorV2\\nfrom albumentations import LongestMaxSize, PadIfNeeded, Compose\\n'"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import os\n",
    "# pip install opencv-python\n",
    "import cv2\n",
    "import torch\n",
    "import numpy as np\n",
    "from torch.utils.data import Dataset, \n",
    "from DataLoader import FLIR_ImagePair_Dataset\n",
    "\"\"\" \n",
    "import albumentations as A\n",
    "# pip install albumentations\n",
    "from albumentations.pytorch import ToTensorV2\n",
    "from albumentations import LongestMaxSize, PadIfNeeded, Compose\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "NUM_CLASSES = 3  # 3 类目标检测\n",
    "SPLIT = [\"train\", \"test\"]  #"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "下面给出一个返回一对图像但是没有融合的加载器，并且这个加载器并不考虑进行数据增强。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Use Case"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "此处我不确定增强的作用是什么，所以这里我选择了False来冒烟测试一遍。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_FLIRPairDataset():\n",
    "    root = \"/root/MyCode/infrared-image-damage-detection-model/dataset/FLIR-align-3class\"\n",
    "    train_ds = FLIRPairDataset(root, split=\"train\", augment=False)\n",
    "    val_ds = FLIRPairDataset(root, split=\"test\", augment=False)\n",
    "    train_loader = DataLoader(\n",
    "        train_ds,\n",
    "        batch_size=4,\n",
    "        shuffle=True,\n",
    "        num_workers=4,\n",
    "        collate_fn=lambda x: tuple(zip(*[(d[\"image\"], d[\"label\"]) for d in x])),\n",
    "    )\n",
    "\n",
    "    val_loader = DataLoader(\n",
    "        val_ds,\n",
    "        batch_size=4,\n",
    "        shuffle=False,\n",
    "        num_workers=4,\n",
    "        collate_fn=lambda x: tuple(zip(*[(d[\"image\"], d[\"label\"]) for d in x])),\n",
    "    )\n",
    "    print(len(train_loader))\n",
    "    print(len(val_loader))  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "无增强加载器测试"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "root = \"/root/MyCode/infrared-image-damage-detection-model/dataset/FLIR-align-3class\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_ds = FLIR_ImagePair_Dataset(root, split=\"train\")\n",
    "val_ds = FLIR_ImagePair_Dataset(root, split=\"test\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_loader = DataLoader(\n",
    "    train_ds,\n",
    "    batch_size=4,\n",
    "    shuffle=True,\n",
    "    num_workers=4,\n",
    "    collate_fn=lambda x: tuple(zip(*[(d[\"image\"], d[\"label\"]) for d in x])),\n",
    ")\n",
    "\n",
    "val_loader = DataLoader(\n",
    "    val_ds,\n",
    "    batch_size=4,\n",
    "    shuffle=False,\n",
    "    num_workers=4,\n",
    "    collate_fn=lambda x: tuple(zip(*[(d[\"image\"], d[\"label\"]) for d in x])),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1033\n",
      "254\n"
     ]
    }
   ],
   "source": [
    "print(len(train_loader))\n",
    "print(len(val_loader))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([3, 640, 640])\n",
      "torch.Size([1, 640, 640])\n",
      "tensor([[1.0000, 0.4500, 0.4570, 0.0156, 0.0156]])\n"
     ]
    }
   ],
   "source": [
    "sample = train_ds[0]\n",
    "print(sample['vis'].shape)   # torch.Size([3, 640, 640])\n",
    "print(sample['ir'].shape)    # torch.Size([1, 640, 640])\n",
    "print(sample['label'])       # (N,5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'DamageDetectionDataset' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "\u001b[1;32m/root/MyCode/infrared-image-damage-detection-model/dataset/FLIR-algin-3class-dataloader.ipynb Cell 16\u001b[0m line \u001b[0;36m<cell line: 2>\u001b[0;34m()\u001b[0m\n\u001b[1;32m      <a href='vscode-notebook-cell://localhost:8080/root/MyCode/infrared-image-damage-detection-model/dataset/FLIR-algin-3class-dataloader.ipynb#X34sdnNjb2RlLXJlbW90ZQ%3D%3D?line=0'>1</a>\u001b[0m \u001b[39m# Step 4: Create train_loader\u001b[39;00m\n\u001b[0;32m----> <a href='vscode-notebook-cell://localhost:8080/root/MyCode/infrared-image-damage-detection-model/dataset/FLIR-algin-3class-dataloader.ipynb#X34sdnNjb2RlLXJlbW90ZQ%3D%3D?line=1'>2</a>\u001b[0m dataset \u001b[39m=\u001b[39m DamageDetectionDataset(image_dir\u001b[39m=\u001b[39m\u001b[39m\"\u001b[39m\u001b[39mpath/to/images\u001b[39m\u001b[39m\"\u001b[39m, annotation_dir\u001b[39m=\u001b[39m\u001b[39m\"\u001b[39m\u001b[39mpath/to/annotations\u001b[39m\u001b[39m\"\u001b[39m)\n\u001b[1;32m      <a href='vscode-notebook-cell://localhost:8080/root/MyCode/infrared-image-damage-detection-model/dataset/FLIR-algin-3class-dataloader.ipynb#X34sdnNjb2RlLXJlbW90ZQ%3D%3D?line=2'>3</a>\u001b[0m train_loader \u001b[39m=\u001b[39m DataLoader(dataset, batch_size\u001b[39m=\u001b[39m\u001b[39m16\u001b[39m, shuffle\u001b[39m=\u001b[39m\u001b[39mTrue\u001b[39;00m, num_workers\u001b[39m=\u001b[39m\u001b[39m4\u001b[39m)\n\u001b[1;32m      <a href='vscode-notebook-cell://localhost:8080/root/MyCode/infrared-image-damage-detection-model/dataset/FLIR-algin-3class-dataloader.ipynb#X34sdnNjb2RlLXJlbW90ZQ%3D%3D?line=4'>5</a>\u001b[0m \u001b[39m# Example usage\u001b[39;00m\n",
      "\u001b[0;31mNameError\u001b[0m: name 'DamageDetectionDataset' is not defined"
     ]
    }
   ],
   "source": [
    "# Step 4: Create train_loader\n",
    "dataset = DamageDetectionDataset(image_dir=\"path/to/images\", annotation_dir=\"path/to/annotations\")\n",
    "train_loader = DataLoader(dataset, batch_size=16, shuffle=True, num_workers=4)\n",
    "\n",
    "# Example usage\n",
    "for batch in train_loader:\n",
    "    images, boxes, labels = batch\n",
    "    print(f\"Batch Images Shape: {images.shape}, Boxes: {boxes}, Labels: {labels}\")\n",
    "    break"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "science39",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
